from tensorflow.keras.models import load_model
from utils.data_generator import test_generator, pred_generator
from utils.image_plot import plot_images
import os

# set gpu
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

# load test set
test_gen = test_generator(
    data_dir='./dataset/natural-scenes/seg_test',
    target_size=(160, 160),
    batch_size=64,
    class_mode='categorical'
)

pred_gen = pred_generator(
    data_dir='./dataset/natural-scenes/seg_test',
    target_size=(160, 160),
    batch_size=64,
    class_mode=None
)

"""
tf.keras.models.load_model
parameter:
- filepath

tf.keras.Sequential.evaluate
parameter:
- x: load val data by ImageDataGenerator

tf.keras.Sequential.predict
parameter:
-x: load test data by ImageDataGenerator
"""

# load model
model_path = './models/model-2020-08-24-15-23-29'
loaded_model = load_model(filepath=model_path)

# metrics for test set
loss, accuracy = loaded_model.evaluate(x=test_gen)

print("loss={}".format(loss))
print("accuracy{}".format(accuracy))

pred_batch = pred_gen.next()
pred_result = loaded_model.predict(x=pred_batch)

class_names = list(test_gen.class_indices.keys())
plot_images(pred_batch, pred_result, class_names)